/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 */

package com.huawei.saashousekeeper.config.dynamicdatasource;

import com.huawei.saashousekeeper.config.TenantContext;
import com.huawei.saashousekeeper.config.balancestrategy.LoadBalanceStrategy;
import com.huawei.saashousekeeper.config.balancestrategy.RandomStrategy;
import com.huawei.saashousekeeper.constants.DbPoolEnum;
import com.huawei.saashousekeeper.constants.LoadBalanceStrategyEnum;
import com.huawei.saashousekeeper.creator.DataSourceCreator;
import com.huawei.saashousekeeper.dbpool.PoolConfig;
import com.huawei.saashousekeeper.exception.RoutingException;
import com.huawei.saashousekeeper.properties.DataSourceBindingProperty;
import com.huawei.saashousekeeper.properties.DataSourceProperty;
import com.huawei.saashousekeeper.properties.DynamicSourceProperties;
import com.huawei.saashousekeeper.properties.MasterSlaveProperty;

import lombok.extern.log4j.Log4j2;

import org.apache.commons.lang.StringUtils;
import org.apache.ibatis.datasource.DataSourceException;
import org.springframework.jdbc.datasource.AbstractDataSource;
import org.springframework.util.CollectionUtils;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

import javax.sql.DataSource;

/**
 * 动态数据源
 *
 * @author lWX1156935
 * @since 2022/4/22
 */
@Log4j2
public class DynamicRoutingDataSource extends AbstractDataSource implements DataSourceRegistry {
    private ConcurrentHashMap<String, DataSourceGroup> groupMap = new ConcurrentHashMap<>(4);

    private ConcurrentHashMap<String, DataSourceBindingProperty> bindingMap = new ConcurrentHashMap<>(4);

    // 负载均衡算法列表
    private Map<String, Class> loadBalanceStrategyMap = null;

    // 数据源创建器列表
    private Map<String, DataSourceCreator> creatorMap = null;

    // 连接池全局配置
    private PoolConfig globalPoolConfig = null;

    private String defaultSource = null;

    /**
     * 构造器
     *
     * @param loadBalanceStrategy 负载均衡算法列表
     * @param dataSourceCreators 数据源创建器列表
     * @param property 配置属性
     */
    public DynamicRoutingDataSource(List<LoadBalanceStrategy> loadBalanceStrategy,
        List<DataSourceCreator> dataSourceCreators, DynamicSourceProperties property) {

        loadBalanceStrategyMap = loadBalanceStrategy.stream()
            .filter(item -> item.getStrategyType() != null)
            .collect(Collectors.toMap(
                item -> org.apache.commons.lang3.StringUtils.defaultString(item.getStrategyType().getCode())
                    .toLowerCase(Locale.ROOT), item -> item.getClass()));
        creatorMap = dataSourceCreators.stream()
            .filter(item -> item.getPoolType() != null)
            .collect(Collectors.toMap(item -> item.getPoolType().getName().toLowerCase(Locale.ROOT), creator -> creator));
        this.globalPoolConfig = (PoolConfig) property;
        Optional.ofNullable(property.getBindingMap()).ifPresent(bindingMap -> this.bindingMap.putAll(bindingMap));
        this.defaultSource = property.getDefaultSource();
        Map<String, MasterSlaveProperty> map = property.getDataSourceMap();
        if (map == null || map.isEmpty()) {
            return;
        }
        for (Map.Entry<String, MasterSlaveProperty> entry : map.entrySet()) {
            registryDataSource(entry.getKey(), entry.getValue());
        }
    }

    @Override
    public void registryDataSource(String sourceName, MasterSlaveProperty masterSlaveProperty) {
        LoadBalanceStrategy loadBalanceStrategy = null;
        try {
            loadBalanceStrategy = (LoadBalanceStrategy) loadBalanceStrategyMap.getOrDefault(
                org.apache.commons.lang3.StringUtils.defaultIfBlank(masterSlaveProperty.getLoadBalanceStrategy(),
                    LoadBalanceStrategyEnum.RANDOM.getCode()).toLowerCase(Locale.ROOT), RandomStrategy.class)
                .getDeclaredConstructor()
                .newInstance();
        } catch (Exception e) {
            throw new DataSourceException(
                "load balance strategy {" + masterSlaveProperty.getLoadBalanceStrategy() + "} init failed", e);
        }

        // 全局连接池配置合并到组级别配置
        masterSlaveProperty.mergeConfig(globalPoolConfig);

        // 解析出区分master|slave|schema的结构
        DataSourceGroup dataSourceGroup = new DataSourceGroup(loadBalanceStrategy,
            createDataSource(masterSlaveProperty.getMaster(), masterSlaveProperty),
            createDataSource(masterSlaveProperty.getSlave(), masterSlaveProperty));
        groupMap.put(sourceName, dataSourceGroup);
    }

    @Override
    public void unRegistryDataSource(String sourceName, List<DataSourceProperty> schemas) {
    }

    @Override
    public DataSourceGroup getDataSourceGroup(String key) {
        String groupName = defaultSource;
        if (bindingMap != null && !bindingMap.isEmpty() && bindingMap.containsKey(key)) {
            groupName = bindingMap.get(key).getGroupName();
        }
        groupName = StringUtils.isBlank(groupName) ? defaultSource : groupName;
        Optional.ofNullable(groupName).orElseThrow(() -> new RoutingException(key + " No Binding Data Source!"));
        DataSourceGroup dataSourceGroup = groupMap.get(groupName);
        Optional.ofNullable(dataSourceGroup)
            .orElseThrow(() -> new RuntimeException(key + " Binding Data Source group not exists"));
        log.warn("{} select DataSource {} success!", key, groupName);
        return dataSourceGroup;
    }

    @Override
    public DataSourceBindingProperty getBidingProperty(String key) {
        if (!bindingMap.containsKey(key)) {
            return null;
        }
        return bindingMap.get(key);
    }

    @Override
    public Connection getConnection() throws SQLException {
        // 获取连接之间切executor.query|update ,已经收集了本次查询是query或者update，此时的分离策略结果包含本次操作
        // 根据传递的租户domain获取所有的租户库，再按分离策略选主从后再走负载均衡选库
        DataSource targetSource = getDataSourceGroup(TenantContext.getDomain()).getDataSource();
        return targetSource.getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        DataSource targetSource = getDataSourceGroup(TenantContext.getDomain()).getDataSource();
        return targetSource.getConnection(username, password);
    }

    private List<DataSource> createDataSource(DataSourceProperty dataSourceProperty, PoolConfig poolConfig) {
        List<DataSourceProperty> list = new ArrayList<>();
        list.add(dataSourceProperty);
        return createDataSource(list, poolConfig);
    }

    private List<DataSource> createDataSource(List<DataSourceProperty> dataSourceProperties, PoolConfig poolConfig) {
        if (CollectionUtils.isEmpty(dataSourceProperties)) {
            return new ArrayList<>();
        }

        // 组级别连接池配置合并到库级别配置
        dataSourceProperties.stream().forEach(item -> {
            item.mergeConfig(poolConfig);
        });

        return dataSourceProperties.stream()
            .map(item -> creatorMap.get(
                org.apache.commons.lang3.StringUtils.defaultIfBlank(item.getPoolName(), DbPoolEnum.DB_DRUID.getName())
                    .toLowerCase(Locale.ROOT)).createDataSource(item))
            .collect(Collectors.toList());
    }
}
